# encoding=utf-8
import argparse
import logging
import time
import torch
from PIL import Image
import cv2 as cv
import os
from EnhanceWrapper import EnhanceWrapper

logging.basicConfig(level=logging.INFO)

# 模型初始化
model = EnhanceWrapper()


def enhance_funds_image(input_path: str, output_path: str):
    if not os.path.exists(input_path):
        logging.error("{} not found, docker will exit".format(input_path))
        return

    logging.info("Now Start funds image...")
    os.makedirs(output_path, exist_ok=True)

    file_list = os.listdir(input_path)
    t1 = time.perf_counter()

    for each_file in file_list:
        t2 = time.perf_counter()
        name, _ = os.path.splitext(each_file)
        png_save_path = os.path.join(output_path, name + ".png")

        with torch.no_grad():
            # 计算
            output = model.infer(os.path.join(input_path, each_file))

            # 输出
            cv.imwrite(png_save_path, output)

        t3 = time.perf_counter()
        logging.info("Done {}, time: {}".format(each_file, t3 - t2))

    t4 = time.perf_counter()
    logging.info("All cost: {}".format(t4 - t1))


def run(base_path: str):
    t1 = time.perf_counter()
    # user_input -> enhance_output
    enhance_funds_image(os.path.join(base_path, "user_input"), os.path.join(base_path, "enhance_output"))

    t2 = time.perf_counter()
    logging.info(
        "All cost {}, {} s/pic".format(t2 - t1, (t2 - t1) / len(os.listdir(os.path.join(base_path, "user_input"))))
    )


if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument("--workdir", type=str)
    args = parser.parse_args()
    if args.workdir == "" or args.workdir is None:
        logging.error("work dir is empty")
        exit(-1)

    run(args.workdir)
